package com.wujunshen.elasticsearch5.support;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.wujunshen.elasticsearch5.wrapper.ESBasicInfo;
import com.wujunshen.elasticsearch5.wrapper.HighLight;
import com.wujunshen.elasticsearch5.wrapper.QueryCondition;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexResponse;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequestBuilder;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.Requests;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: frankwoo(吴峻申) <br>
 * @date: 2018/7/2 <br>
 * @time: 15:25 <br>
 * @mail: frank_wjs@hotmail.com <br>
 */
@Slf4j
public class Elasticsearch5Template {
    @Resource
    private TransportClient esClient;

    private ObjectMapper mapper = new ObjectMapper();

    /**
     * 创建索引，相当于数据库创建schema，类似于DDL脚本
     *
     * @param index   索引，类似数据库
     * @param type    类型，类似表
     * @param mapping 具体mapping
     * @return 操作是否成功标志
     * @throws IOException 异常
     */
    public boolean createMapping(String index, String type, XContentBuilder mapping)
            throws IOException {
        log.info("mapping is:{}", mapping.string());

        PutMappingRequest mappingRequest = Requests.
                putMappingRequest(index).source(mapping).type(type);
        PutMappingResponse putMappingResponse =
                esClient.admin().indices().putMapping(mappingRequest).actionGet();
        return putMappingResponse.isAcknowledged();
    }

    /**
     * 获取mapping
     *
     * @param index 索引，类似数据库
     * @param type  类型，类似表
     * @return mapping的字符串格式
     * @throws IOException 异常
     */
    public String getMapping(String index, String type) throws IOException {
        ImmutableOpenMap<String, MappingMetaData> mappings = esClient.admin().cluster().prepareState().execute()
                .actionGet().getState().getMetaData().getIndices().get(index).getMappings();

        return mappings.get(type).source().string();
    }

    /**
     * 获取所有mapping
     *
     * @param index 索引，类似数据库
     * @return 所有mapping
     * @throws IOException 异常
     */
    public List<Map<String, Object>> getAllMapping(String index) throws IOException {
        List<Map<String, Object>> result = new ArrayList<>();
        ImmutableOpenMap<String, MappingMetaData> mappings = esClient.admin().cluster().prepareState().execute()
                .actionGet().getState().getMetaData().getIndices().get(index).getMappings();

        for (ObjectObjectCursor<String, MappingMetaData> cursor : mappings) {
            log.info("type is:{}", cursor.key); // 索引下的每个type
            result.add(cursor.value.getSourceAsMap());
        }
        return result;
    }

    /**
     * 创建索引
     *
     * @param index 索引，类似数据库
     * @return 操作是否成功标志
     */
    public boolean createIndex(String index) {
        if (isExistedIndex(index)) {
            return false;
        }
        CreateIndexResponse indexResponse = esClient.admin().indices().prepareCreate(index).get();

        return indexResponse.isAcknowledged(); // true表示创建成功
    }

    /**
     * 删除索引
     *
     * @param index 索引，类似数据库
     * @return 操作是否成功标志
     */
    public boolean deleteIndex(String index) {
        if (!isExistedIndex(index)) {
            return false;
        }
        DeleteIndexResponse deleteIndexResponse = esClient.admin().indices().prepareDelete(index).execute().actionGet();
        return deleteIndexResponse.isAcknowledged(); // true表示创建成功
    }

    /**
     * 判断索引是否存在
     *
     * @param index 索引，类似数据库
     * @return 操作是否成功标志
     */
    public boolean isExistedIndex(String index) {
        return esClient.admin().indices().prepareExists(index).execute().actionGet().isExists();
    }

    /**
     * 添加数据
     *
     * @param esBasicInfo es基本信息
     * @param object      数据序列化对象
     * @return 操作是否成功标志
     * @throws IOException 异常
     */
    public boolean addData(ESBasicInfo esBasicInfo, Object object)
            throws IOException {
        IndexResponse result = esClient.prepareIndex(esBasicInfo.getIndex(),
                esBasicInfo.getType(), esBasicInfo.getIds()[0])
                .setSource(mapper.writeValueAsString(object), XContentType.JSON)
                .get();

        return result.status().getStatus() == 201;
    }

    /**
     * 更新数据
     *
     * @param esBasicInfo es基本信息
     * @param object      数据序列化对象
     * @return 操作是否成功标志
     * @throws IOException 异常
     */
    public boolean updateData(ESBasicInfo esBasicInfo, Object object)
            throws IOException {
        UpdateResponse result = esClient.prepareUpdate(esBasicInfo.getIndex(),
                esBasicInfo.getType(), esBasicInfo.getIds()[0])
                .setDoc(mapper.writeValueAsString(object), XContentType.JSON)
                .get();

        return result.status().getStatus() == 200;
    }

    /**
     * 删除数据
     *
     * @param esBasicInfo es基本信息
     * @return 操作是否成功标志
     */
    public boolean deleteData(ESBasicInfo esBasicInfo) {
        DeleteResponse result = esClient.prepareDelete(esBasicInfo.getIndex(),
                esBasicInfo.getType(), esBasicInfo.getIds()[0]).get();

        return result.status().getStatus() == 200;
    }

    /**
     * 批量添加数据
     *
     * @param esBasicInfo es基本信息
     * @param object      数据序列化对象
     * @return 批量添加的数据条数
     * @throws IOException 异常
     */
    public int addBatchData(ESBasicInfo esBasicInfo, Object object)
            throws IOException {
        BulkRequestBuilder bulkRequest = esClient.prepareBulk();

        for (String id : esBasicInfo.getIds()) {
            bulkRequest.add(esClient.prepareIndex(esBasicInfo.getIndex(),
                    esBasicInfo.getType(), id).
                    setSource(mapper.writeValueAsString(object), XContentType.JSON));
        }
        bulkRequest.execute().actionGet();

        return bulkRequest.numberOfActions();
    }

    /**
     * 批量更新数据
     *
     * @param esBasicInfo es基本信息
     * @param object      数据序列化对象
     * @return 批量更新的数据条数
     * @throws IOException 异常
     */
    public int updateBatchData(ESBasicInfo esBasicInfo, Object object)
            throws IOException {
        BulkRequestBuilder bulkRequest = esClient.prepareBulk();

        for (String id : esBasicInfo.getIds()) {
            bulkRequest.add(esClient.prepareUpdate(esBasicInfo.getIndex(),
                    esBasicInfo.getType(), id).
                    setDoc(mapper.writeValueAsString(object), XContentType.JSON));
        }

        bulkRequest.execute().actionGet();

        return bulkRequest.numberOfActions();
    }

    /**
     * 批量删除数据
     *
     * @param esBasicInfo es基本信息
     * @return 批量删除的数据条数
     */
    public int deleteBatchData(ESBasicInfo esBasicInfo) {
        BulkRequestBuilder bulkRequest = esClient.prepareBulk();

        for (String id : esBasicInfo.getIds()) {
            bulkRequest.add(esClient.prepareDelete(esBasicInfo.getIndex(),
                    esBasicInfo.getType(), id));
        }

        BulkResponse response = bulkRequest.execute().actionGet();
        log.info("status is:{}", response.status().getStatus());

        return bulkRequest.numberOfActions();
    }

    /**
     * 通过ID查询数据
     *
     * @param esBasicInfo es基本信息
     * @param clazz       要序列化的查询结果对象类型
     * @param <T>         查询结果对象具体类型
     * @return 查询结果对象
     * @throws IOException 异常
     */
    public <T> T query(ESBasicInfo esBasicInfo, Class<T> clazz) throws IOException {
        GetRequestBuilder requestBuilder = esClient.prepareGet(
                esBasicInfo.getIndex(), esBasicInfo.getType(), esBasicInfo.getIds()[0]);
        GetResponse response = requestBuilder.execute().actionGet();

        return response.getSourceAsString() != null ? mapper.readValue(response.getSourceAsString(), clazz) : null;
    }

    /**
     * 执行查询的封装方法
     *
     * @param index          索引，类似数据库
     * @param queryCondition 查询条件类
     * @param type           类型，类似表
     * @return SearchResponse对象
     */
    public SearchResponse executeQuery(String index,
                                       QueryCondition queryCondition,
                                       String... type) {
        return esClient.prepareSearch(index)
                .setTypes(type)
                .setSearchType(queryCondition.getSearchType())
                .setScroll(new TimeValue(queryCondition.getMillis()))
                .setQuery(queryCondition.getQueryBuilder())
                .setSize(queryCondition.getSize()).execute().actionGet();
    }

    /**
     * 解析SearchResponse类
     *
     * @param clazz    要序列化的查询结果对象类型
     * @param response 要解析的SearchResponse类
     * @param <T>      查询结果对象具体类型
     * @return 查询结果对象
     * @throws IOException 异常
     */
    public <T> List<T> analyzeSearchResponse(Class<T> clazz, SearchResponse response) throws IOException {
        SearchHits searchHits = response.getHits();

        List<T> result = new ArrayList<>();
        for (SearchHit hit : searchHits) {
            result.add(mapper.readValue(hit.getSourceAsString(), clazz));
        }
        return result;
    }

    /**
     * 高亮结果集 特殊处理
     *
     * @param esBasicInfo    es基本信息
     * @param queryCondition 查询条件类
     * @param highLight      高亮设置类
     * @return 高亮结果集
     */
    public List<Map<String, Object>> highLightResultSet(ESBasicInfo esBasicInfo,
                                                        QueryCondition queryCondition,
                                                        HighLight highLight) {
        SearchResponse response = esClient.prepareSearch(esBasicInfo.getIndex())
                .setTypes(esBasicInfo.getType())
                .setSearchType(queryCondition.getSearchType())
                .setScroll(new TimeValue(queryCondition.getMillis()))
                .setQuery(queryCondition.getQueryBuilder())
                .setSize(queryCondition.getSize())
                .highlighter(highLight.getBuilder())//添加高亮的字段
                .execute().actionGet();

        String highlightField = highLight.getField();
        List<Map<String, Object>> sourceList = new ArrayList<>();

        for (SearchHit searchHit : response.getHits()) {
            Map<String, Object> element = new HashMap<>();
            StringBuilder stringBuilder = new StringBuilder();
            if (StringUtils.isNotEmpty(highlightField)) {
                Text[] text = searchHit.getHighlightFields().get(highlightField).getFragments();

                if (text != null) {
                    for (Text str : text) {
                        stringBuilder.append(str.string());
                    }
                    //遍历 高亮结果集
                    log.info("遍历 高亮结果集{}", stringBuilder.toString());
                    element.put(highlightField, stringBuilder.toString());
                }
            }
            sourceList.add(element);
        }

        return sourceList;
    }
}